import math
from abc import ABC, abstractmethod
from typing import List, Dict, Optional, Sequence
from dalchemy.prompts import PromptTemplate
from dalchemy.parsers import DalchemyParser
from dalchemy.data import TextHelper

txthelper = TextHelper()


class DalchemyData(object):
    ''' 数据基类：负责加载数据、处理数据、加载模板、格式化数据 '''
    ''' TODO: 接受一个file_path、loader、prompt，然后写一个生成prompt '''
    fields: List[str] = []
    field_format: str = "format_info"

    # field_example: str = "example"

    def __init__(self, data_file, template_file):
        self.data_file = data_file
        self.template_file = template_file
        # self.input: List[str] = []
        # self.prompt: str = ""

    # @abstractmethod
    def fetch_raw_data(self) -> List:
        '''读取原始数据，默认加载json数据'''
        raise NotImplemented

    # @abstractmethod
    def process_data(self, data_ls: List):
        ''' 读取、处理原始数据，构造需要格式的json，如dict_ls = [{"question":"q1", "answer":"a1"},] '''
        raise NotImplemented

    def format_template(self, dict_ls: List[Dict[str, str]],
                        parser: Optional[DalchemyParser] = None):
        ''' 根据模板和输入，产生prompt '''
        # 1.read template
        template = txthelper.read_file(self.template_file)
        # 2.init PromptTemplate
        try:
            format_info = {self.field_format: parser.get_format_instructions()}
            self.prompt = PromptTemplate(template=template,
                                         input_variables=self.fields,
                                         partial_variables=format_info)

        except Exception as e:
            self.prompt = PromptTemplate(template=template,
                                         input_variables=self.fields)


        # 3.format data
        prompt_ls = [self.prompt.format(**dic) for dic in dict_ls]
        return prompt_ls

    def generate_prompts(self, parser: Optional[DalchemyParser] = None)->List[str]:
        ''' 暴露给外部的接口，一键完成数据加载、处理、检查、格式化模板 '''
        data_ls = self.fetch_raw_data() # fetch可以写在process_datane内部
        dict_ls = self.process_data(data_ls) # 字典列表
        self.dict_ls = dict_ls
        self.validate_data_fields(dict_ls) # 检验dict字段是否符合类属性fields
        prompts = self.format_template(dict_ls, parser)
        return prompts

    @classmethod
    def validate_data_fields(cls, dict_ls: List[Dict[str, str]]):
        '''检查数据中的字段是否符合定义的字段信息'''
        for dic in dict_ls:
            for field_name in cls.fields:
                if field_name not in dic.keys():
                    raise ValueError(f"Field '{field_name}' is missing in data.")

    @classmethod
    def make_batches(cls,
                    prompts: Sequence[str],
                    batch_size: Optional[int] = 1):
        prompts = [p.strip() for p in prompts]
        n_examples = len(prompts)
        n_batches = int(math.ceil(n_examples / batch_size))
        prompt_batches = [prompts[batch_id * batch_size: (batch_id + 1) * batch_size]
                          for batch_id in range(n_batches)]
        return prompt_batches

    # @property
    # @abstractmethod
    # def identifying_params(self):
    #     """返回模板需要的参数."""
    #     return {}
